#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
Created on Wed Nov  2 09:14:04 2016

@author: kevin
"""
import tensorflow as tf
import numpy as np
import cv2
import os

from utils import weight_variable, bias_variable, conv2d
#import matlab.engine
#eng = matlab.engine.start_matlab()

class CNN1(object):
     """
     input:  originalImage                       [batch_sz, im_h, im_w, im_ch]
     output: feature vector to initialize rnn    [batch_sz, rnn_cell_sz]
     """
     def __init__(self, config, image_ph):
          self.batch_sz = config.batch_sz
          self.im_sz = config.img_sz
          self.im_ch = config.img_channels

          self.output_sz = config.rnn_input_sz
          self.image_ph = image_ph
          self.init_weights()
          
     def init_weights(self):
          # first hidden layer
          self.w_conv1 = weight_variable([5,5,self.im_ch,32])
          self.b_conv1 = bias_variable([32])
          # second hidden layer
          self.w_conv2 = weight_variable([5,5,32,64])
          self.b_conv2 = bias_variable([64])
          # first densely connected layer
          self.w_fc1 = weight_variable([self.im_sz*self.im_sz*64, 1024])
          self.b_fc1 = bias_variable([1024])
          # second densely connected layer --- 1000 dim vector
          self.w_fc2 = weight_variable([1024,self.output_sz])
          self.b_fc2 = bias_variable([self.output_sz])
          
     def __call__(self,):
          
          X_image = tf.reshape(self.image_ph,[-1,self.im_sz,self.im_sz,self.im_ch])

          h_conv1 = tf.nn.relu(conv2d(X_image, self.w_conv1)+self.b_conv1)
          h_conv2 = tf.nn.relu(conv2d(h_conv1, self.w_conv2)+self.b_conv2)
          
          h_conv2_flat = tf.reshape(h_conv2, [-1, self.im_sz*self.im_sz*64])
          
          h_fc1 = tf.nn.relu(tf.matmul(h_conv2_flat, self.w_fc1) + self.b_fc1)
          # fix CNN1 
          h_conv1 = tf.stop_gradient(h_conv1)
          h_conv2 = tf.stop_gradient(h_conv2)
          h_conv2_flat = tf.stop_gradient(h_conv2_flat)
          h_fc1 = tf.stop_gradient(h_fc1)
        
          y = tf.nn.softmax(tf.matmul(h_fc1, self.w_fc2) + self.b_fc2)
          y = tf.stop_gradient(y)
          return y

        
class CNN2(object):
     """
     input:    originalImage && accumulated strokeMap    [batch_sz, im_h, im_w, im_ch]
     output:   alphaMatte                                [batch_sz, im_h, im_w, 1]
     """
     def __init__(self, config):
          
          self.batch_sz = config.batch_sz
          self.im_sz = config.img_sz
          self.im_ch = config.img_channels
          
          self.output_sz = config.img_sz
          self.output_ch = config.alphaMatte_channels
          
          
     def op(self, image_ph, strokeMap_ph):
          self.image_ph = image_ph
          self.strokeMap_ph = strokeMap_ph
          
          X_image = tf.reshape(self.image_ph, [-1, self.im_sz, self.im_sz, self.im_ch])
          X_stroke = tf.reshape(self.strokeMap_ph, [-1, self.im_sz, self.im_sz, self.im_ch])
#          X_image = self.image_ph
#          X_stroke = self.strokeMap_ph
          
          counter = tf.shape(X_image)[0]
#          with tf.Session() as sess:
#               counter_val = sess.run(counter)
#               X_image_val = sess.run(X_image).tolist()
#               X_stroke_val = sess.run(X_stroke).tolist()
          with tf.Session() as sess:
               counter_val,X_image_val,X_stroke_val = sess.run([counter,X_image,X_stroke])
               
          X_image_val = X_image_val.tolist()
          X_stroke_val = X_stroke_val.tolist()
          mattes = np.random.random((counter_val,self.im_sz, self.im_sz, 1))  # only for initialize
          for i in range(counter_val):
            X_image_i = np.array(X_image_val[i])
            X_image_i = np.uint8(X_image_i)
            #X_image_i = cv2.cvtColor(X_image_i,cv2.COLOR_BGR2RGB)
            cv2.imwrite('./first_try_DATA/tmp/image.png',X_image_i)
            #X_image_i=Image.fromarray(X_image_i)
            #X_image_i.convert("RGB")
            #X_image_i.save('./first_try_DATA/tmp/image.png')
               
            X_stroke_i = np.array(X_stroke_val[i])
            X_stroke_i = np.uint8(X_stroke_i)
            cv2.imwrite('./first_try_DATA/tmp/stroke.png',X_stroke_i)
            
            systemCall='./SharedMatting -i ./first_try_DATA/tmp/image.png -t ./first_try_DATA/tmp/stroke.png -a'
            os.system(systemCall)
            matte=cv2.imread('shared_matting_alpha.bmp')
            os.system('rm shared_matting_alpha.bmp')
            cv2.imwrite('./first_try_DATA/tmp/matte.png',matte)
            matte=cv2.imread('./first_try_DATA/tmp/matte.png',0)
            matte=matte.reshape((self.im_sz,self.im_sz,self.output_ch))
            #cv2.imshow('Before',matte)
            #cv2.waitKey(0)
            mattes[i]=matte
            #cv2.imshow('After',mattes[i])
            #cv2.waitKey(0)
            #cv2.destroyAllWindows()
          return mattes
          
          
          
          
          
          
          
          
          
          
          
        
        
        
